from  schemas.agent_schemas import Task
from app.agent.manus import Manus
from app.logger import logger
from datetime import datetime
import uuid
import asyncio

class TaskManager:
    def __init__(self):
        self.tasks = {}
        self.queues = {}

    def create_task(self, prompt: str) -> Task:
        task_id = str(uuid.uuid4())
        task = Task(
            id=task_id, prompt=prompt, created_at=datetime.now(), status="waiting"
        )
        self.tasks[task_id] = task
        self.queues[task_id] = asyncio.Queue()
        return task

    async def update_task_step(
        self, task_id: str, step: int, result: str, step_type: str = "step"
    ):
        if task_id in self.tasks:
            task = self.tasks[task_id]
            task.steps.append({"step": step, "result": result, "type": step_type})
            await self.queues[task_id].put(
                {"type": step_type, "step": step, "result": result}
            )
            await self.queues[task_id].put(
                {"type": "status", "status": task.status, "steps": task.steps}
            )

    async def complete_task(self, task_id: str):
        if task_id in self.tasks:
            task = self.tasks[task_id]
            task.status = "completed"
            await self.queues[task_id].put(
                {"type": "status", "status": task.status, "steps": task.steps}
            )
            await self.queues[task_id].put({"type": "complete"})

    async def fail_task(self, task_id: str, error: str):
        if task_id in self.tasks:
            self.tasks[task_id].status = f"failed: {error}"
            await self.queues[task_id].put({"type": "error", "message": error})

    async def run_task(self, task_id: str, prompt: str):
        try:
            self.tasks[task_id].status = "running"

            agent = Manus(
                name="Manus",
                description="A versatile agent that can solve various tasks using multiple tools",
            )

            async def on_think(thought):
                await self.update_task_step(task_id, 0, thought, "think")

            async def on_tool_execute(tool, input):
                await self.update_task_step(
                    task_id, 0, f"Executing tool: {tool}\nInput: {input}", "tool"
                )

            async def on_action(action):
                await self.update_task_step(
                    task_id, 0, f"Executing action: {action}", "act"
                )

            async def on_run(step, result):
                await self.update_task_step(task_id, step, result, "run")

            class SSELogHandler:
                def __init__(self, task_manager, task_id):
                    self.task_manager = task_manager
                    self.task_id = task_id

                async def __call__(self, message):
                    import re
                    cleaned_message = re.sub(r"^.*? - ", "", message)

                    event_type = "log"
                    if "✨ Manus's thoughts:" in cleaned_message:
                        event_type = "think"
                    elif "🛠️ Manus selected" in cleaned_message:
                        event_type = "tool"
                    elif "🎯 Tool" in cleaned_message:
                        event_type = "act"
                    elif "📝 Oops!" in cleaned_message:
                        event_type = "error"
                    elif "🏁 Special tool" in cleaned_message:
                        event_type = "complete"

                    await self.task_manager.update_task_step(
                        self.task_id, 0, cleaned_message, event_type
                    )

            sse_handler = SSELogHandler(self, task_id)
            logger.add(sse_handler)

            result = await agent.run(prompt)
            await self.update_task_step(task_id, 1, result, "result")
            await self.complete_task(task_id)
        except Exception as e:
            await self.fail_task(task_id, str(e))

# 创建全局TaskManager实例
task_manager = TaskManager()

